#!/usr/bin/python3
# coding=utf-8

import torch
import torch.nn.functional as F
import numpy as np
import os

def im2col_torch(x, kernel_size, stride=1, padding=0):
    """PyTorch实现的im2col函数
    
    参数:
        x: 输入张量，形状为 [N, C, H, W]
        kernel_size: 卷积核大小，可以是int或tuple
        stride: 步长
        padding: 填充大小
        
    返回:
        im2col展开后的矩阵，形状为 [N, C*kH*kW, L]
    """
    if isinstance(kernel_size, int):
        kernel_size = (kernel_size, kernel_size)
    
    # 使用PyTorch的unfold函数实现im2col
    col = F.unfold(x, kernel_size, stride=stride, padding=padding)
    return col

def gen_golden_data_im2col():
    """生成im2col操作的测试数据"""
    # 配置参数
    dtype = torch.float32
    input_shape = [9, 10, 23, 1971]  # [N, C, H, W]
    
    # 卷积参数
    kernel_size = (4, 3)  # 卷积核大小
    stride = 4            # 步长
    padding = 8           # 填充大小
    
    # 生成输入数据 - 使用有规律的数据便于验证
    torch.manual_seed(42)  # 固定随机种子以便复现
    # 使用0-1之间的随机数，便于观察
    x = torch.rand(input_shape, dtype=dtype)
    
    # 创建输入输出目录
    os.makedirs("./input", exist_ok=True)
    os.makedirs("./output", exist_ok=True)
    
    # 计算im2col结果
    golden = im2col_torch(x, kernel_size, stride, padding)
    
    # 转换为numpy用于保存
    x_np = x.numpy()
    golden_np = golden.numpy()
    
    # 保存二进制文件
    x_np.tofile("./input/input_x.bin")
    golden_np.tofile("./output/golden.bin")

if __name__ == "__main__":
    gen_golden_data_im2col()